from typing import Tuple

import numpy as np
import paddle

from ppsci.geometry import geometry
from ppsci.geometry import sampler


class Cylinder3D(geometry.Geometry):
    """Class for Cylinder (3D)

    Args:
        center (Tuple[float, float, float]): Center point of the cylinder base [x0, y0, z0].
        radius (float): Radius of the cylinder.
        height (float): Height of the cylinder.

    Examples:
        >>> import ppsci
        >>> geom = ppsci.geometry.Cylinder((0.0, 0.0, 0.0), 1.0, 2.0)
    """

    def __init__(
        self, center: Tuple[float, float, float], radius: float, height: float
    ):
        self.center = np.array(center, dtype=paddle.get_default_dtype())
        self.radius = radius
        self.height = height

        # Bounding box: [min_x, min_y, min_z] to [max_x, max_y, max_z]
        bbox_min = self.center - np.array([radius, radius, 0])
        bbox_max = self.center + np.array([radius, radius, height])
        diam = np.linalg.norm(bbox_max - bbox_min)

        super().__init__(3, (bbox_min, bbox_max), diam)
        self.base_area = np.pi * radius ** 2
        self.lateral_area = 2 * np.pi * radius * height
        self.surface_area = 2 * self.base_area + self.lateral_area
        self.volume = self.base_area * height

    def is_inside(self, x: np.ndarray) -> np.ndarray:
        """Check if points are inside the cylinder"""
        # Project point to base plane
        xy_dist = np.linalg.norm(x[:, :2] - self.center[:2], axis=1)
        # Check if within z bounds
        z_inside = (x[:, 2] >= self.center[2]) & (
            x[:, 2] <= self.center[2] + self.height
        )
        return (xy_dist <= self.radius) & z_inside

    def on_boundary(self, x: np.ndarray) -> np.ndarray:
        """Check if points are on the cylinder boundary"""
        xy_dist = np.linalg.norm(x[:, :2] - self.center[:2], axis=1)
        atol = 1e-6

        # Check base (z = center_z)
        on_base = np.isclose(x[:, 2], self.center[2], atol=atol) & (
            xy_dist <= self.radius + atol
        )
        # Check top (z = center_z + height)
        on_top = np.isclose(x[:, 2], self.center[2] + self.height, atol=atol) & (
            xy_dist <= self.radius + atol
        )
        # Check lateral surface
        on_side = (
            np.isclose(xy_dist, self.radius, atol=atol)
            & (x[:, 2] >= self.center[2] - atol)
            & (x[:, 2] <= self.center[2] + self.height + atol)
        )

        return on_base | on_top | on_side

    def boundary_normal(self, x: np.ndarray) -> np.ndarray:
        """Compute the unit normal at x for the cylinder boundary"""
        normals = np.zeros_like(x)
        xy_dist = np.linalg.norm(x[:, :2] - self.center[:2], axis=1)
        atol = 1e-6

        # Base (z = center_z)
        base_mask = np.isclose(x[:, 2], self.center[2], atol=atol) & (
            xy_dist <= self.radius + atol
        )
        normals[base_mask, 2] = -1.0  # Downward normal

        # Top (z = center_z + height)
        top_mask = np.isclose(x[:, 2], self.center[2] + self.height, atol=atol) & (
            xy_dist <= self.radius + atol
        )
        normals[top_mask, 2] = 1.0  # Upward normal

        # Lateral surface
        side_mask = (
            np.isclose(xy_dist, self.radius, atol=atol)
            & (x[:, 2] > self.center[2] + atol)
            & (x[:, 2] < self.center[2] + self.height - atol)
        )

        if np.any(side_mask):
            # Compute radial normal
            dx = x[side_mask, 0] - self.center[0]
            dy = x[side_mask, 1] - self.center[1]
            dist = np.sqrt(dx ** 2 + dy ** 2)
            # Avoid division by zero
            dist[dist < atol] = 1.0
            normals[side_mask, 0] = dx / dist
            normals[side_mask, 1] = dy / dist

        return normals

    def random_points(self, n: int, random: str = "pseudo") -> np.ndarray:
        """Generate random points inside the cylinder"""
        # Generate points in cylindrical coordinates [r, theta, z]
        r = (
            np.sqrt(sampler.sample(n, 1, random)) * self.radius
        )  # sqrt for uniform distribution
        theta = 2 * np.pi * sampler.sample(n, 1, random)
        z = sampler.sample(n, 1, random) * self.height + self.center[2]

        # Convert to Cartesian coordinates
        x = self.center[0] + r * np.cos(theta)
        y = self.center[1] + r * np.sin(theta)
        return np.hstack((x, y, z))

    def uniform_points(self, n: int, boundary: bool = True) -> np.ndarray:
        """Generate uniform points inside the cylinder"""
        # Estimate grid size based on desired point density
        step = (self.volume / n) ** (1 / 3)

        # Determine grid sizes
        nr = max(1, int(np.ceil(self.radius / step)))
        ntheta = max(3, int(np.ceil(2 * np.pi * self.radius / step)))
        nz = max(1, int(np.ceil(self.height / step)))

        # Generate cylindrical grid
        r = np.linspace(0, self.radius, nr, dtype=paddle.get_default_dtype())
        theta = np.linspace(
            0, 2 * np.pi, ntheta, endpoint=False, dtype=paddle.get_default_dtype()
        )
        z = np.linspace(
            self.center[2],
            self.center[2] + self.height,
            nz,
            dtype=paddle.get_default_dtype(),
        )

        # Create points
        points = []
        for r_val in r:
            for theta_val in theta:
                for z_val in z:
                    # Skip duplicate center points
                    if r_val == 0 and theta_val > 0:
                        continue
                    points.append(
                        [
                            self.center[0] + r_val * np.cos(theta_val),
                            self.center[1] + r_val * np.sin(theta_val),
                            z_val,
                        ]
                    )

        points = np.array(points, dtype=paddle.get_default_dtype())
        return points[:n] if len(points) > n else points

    def random_boundary_points(self, n: int, random: str = "pseudo") -> np.ndarray:
        """Generate random points on the cylinder boundary"""
        # Calculate point distribution based on surface area
        n_base = max(1, int(n * self.base_area / self.surface_area))
        n_top = max(1, int(n * self.base_area / self.surface_area))
        n_side = n - n_base - n_top

        points = []

        # Base points
        if n_base > 0:
            r_base = np.sqrt(sampler.sample(n_base, 1, random)) * self.radius
            theta_base = 2 * np.pi * sampler.sample(n_base, 1, random)
            x_base = self.center[0] + r_base * np.cos(theta_base)
            y_base = self.center[1] + r_base * np.sin(theta_base)
            z_base = np.full((n_base, 1), self.center[2])
            base_points = np.hstack((x_base, y_base, z_base))
            points.append(base_points)

        # Top points
        if n_top > 0:
            r_top = np.sqrt(sampler.sample(n_top, 1, random)) * self.radius
            theta_top = 2 * np.pi * sampler.sample(n_top, 1, random)
            x_top = self.center[0] + r_top * np.cos(theta_top)
            y_top = self.center[1] + r_top * np.sin(theta_top)
            z_top = np.full((n_top, 1), self.center[2] + self.height)
            top_points = np.hstack((x_top, y_top, z_top))
            points.append(top_points)

        # Side points
        if n_side > 0:
            theta_side = 2 * np.pi * sampler.sample(n_side, 1, random)
            z_side = sampler.sample(n_side, 1, random) * self.height + self.center[2]
            x_side = self.center[0] + self.radius * np.cos(theta_side)
            y_side = self.center[1] + self.radius * np.sin(theta_side)
            side_points = np.hstack((x_side, y_side, z_side))
            points.append(side_points)

        return np.vstack(points)

    def uniform_boundary_points(self, n: int) -> np.ndarray:
        """Generate uniform points on the cylinder boundary"""
        # Calculate surface density
        density = n / self.surface_area

        # Lateral surface grid
        ntheta = max(3, int(np.ceil(density * 2 * np.pi * self.radius)))
        nz = max(1, int(np.ceil(density * self.height)))
        lateral_points = []
        theta = np.linspace(
            0, 2 * np.pi, ntheta, endpoint=False, dtype=paddle.get_default_dtype()
        )
        z = np.linspace(
            self.center[2],
            self.center[2] + self.height,
            nz,
            dtype=paddle.get_default_dtype(),
        )
        for theta_val in theta:
            for z_val in z:
                lateral_points.append(
                    [
                        self.center[0] + self.radius * np.cos(theta_val),
                        self.center[1] + self.radius * np.sin(theta_val),
                        z_val,
                    ]
                )

        # Cap surfaces grid
        cap_points = []
        nr = max(1, int(np.ceil(density * self.radius)))
        ntheta_cap = max(3, int(np.ceil(density * 2 * np.pi * self.radius)))
        r = np.linspace(0, self.radius, nr, dtype=paddle.get_default_dtype())
        theta = np.linspace(
            0, 2 * np.pi, ntheta_cap, endpoint=False, dtype=paddle.get_default_dtype()
        )

        # Base cap
        for r_val in r:
            for theta_val in theta:
                if r_val == 0 and theta_val > 0:
                    continue
                cap_points.append(
                    [
                        self.center[0] + r_val * np.cos(theta_val),
                        self.center[1] + r_val * np.sin(theta_val),
                        self.center[2],  # base z
                    ]
                )

        # Top cap
        for r_val in r:
            for theta_val in theta:
                if r_val == 0 and theta_val > 0:
                    continue
                cap_points.append(
                    [
                        self.center[0] + r_val * np.cos(theta_val),
                        self.center[1] + r_val * np.sin(theta_val),
                        self.center[2] + self.height,  # top z
                    ]
                )

        # Combine points
        points = np.vstack(
            [
                np.array(lateral_points, dtype=paddle.get_default_dtype()),
                np.array(cap_points, dtype=paddle.get_default_dtype()),
            ]
        )
        return points[:n] if len(points) > n else points

    def sdf_func(self, points: np.ndarray) -> np.ndarray:
        """Compute signed distance field for cylinder.

        Args:
            points (np.ndarray): Points to evaluate (shape [N, 3])

        Returns:
            np.ndarray: SDF values (shape [N, 1])
        """
        if points.shape[1] != 3:
            raise ValueError(f"Points must be 3D, got shape {points.shape}")

        # Vector from center to points in xy-plane
        xy_vec = points[:, :2] - self.center[:2]
        xy_dist = np.linalg.norm(xy_vec, axis=1)

        # Distance to side surface
        side_dist = xy_dist - self.radius

        # Distance to bases
        z_dist_bottom = points[:, 2] - self.center[2]
        z_dist_top = (self.center[2] + self.height) - points[:, 2]

        # For points inside cylinder extent
        mask_inside_z = (points[:, 2] > self.center[2]) & (
            points[:, 2] < self.center[2] + self.height
        )

        # For points outside cylinder extent
        dist_z = np.minimum(z_dist_bottom, z_dist_top)
        dist_xy = np.maximum(side_dist, -dist_z)

        # Combine results
        sdf = np.where(
            mask_inside_z,
            side_dist,
            np.sqrt(np.maximum(side_dist, 0) ** 2 + np.minimum(dist_z, 0) ** 2),
        )

        return sdf.reshape(-1, 1)
